/*
 * SPI函数的接口
 *
 * src/pl/plpython/plpy_spi.c
 */

#include "postgres.h"

#include <limits.h>

#include "access/htup_details.h"
#include "access/xact.h"
#include "catalog/pg_type.h"
#include "executor/spi.h"
#include "mb/pg_wchar.h"
#include "parser/parse_type.h"
#include "plpy_elog.h"
#include "plpy_main.h"
#include "plpy_planobject.h"
#include "plpy_plpymodule.h"
#include "plpy_procedure.h"
#include "plpy_resultobject.h"
#include "plpy_spi.h"
#include "plpython.h"
#include "utils/memutils.h"
#include "utils/syscache.h"

static PyObject *fc_PLy_spi_execute_query(char *fc_query, long fc_limit);
static PyObject *fc_PLy_spi_execute_fetch_result(SPITupleTable *fc_tuptable,
											  uint64 fc_rows, int fc_status);
static void fc_PLy_spi_exception_set(PyObject *fc_excclass, ErrorData *fc_edata);


/* 准备(query="select * from foo")
 * 准备(query="select * from foo where bar = $1", params=["text"])
 * 准备(query="select * from foo where bar = $1", params=["text"], limit=5)
 */
PyObject * PLy_spi_prepare(PyObject *fc_self, PyObject *fc_args)
{
	PLyPlanObject *fc_plan;
	PyObject   *fc_list = NULL;
	PyObject   *volatile fc_optr = NULL;
	char	   *fc_query;
	PLyExecutionContext *fc_exec_ctx = PLy_current_execution_context();
	volatile MemoryContext fc_oldcontext;
	volatile ResourceOwner fc_oldowner;
	volatile int fc_nargs;

	if (!PyArg_ParseTuple(fc_args, "s|O:prepare", &fc_query, &fc_list))
		return NULL;

	if (fc_list && (!PySequence_Check(fc_list)))
	{
		PLy_exception_set(PyExc_TypeError,
						  "second argument of plpy.prepare must be a sequence");
		return NULL;
	}

	if ((fc_plan = (PLyPlanObject *) PLy_plan_new()) == NULL)
		return NULL;

	fc_plan->mcxt = AllocSetContextCreate(TopMemoryContext,
									   "PL/Python plan context",
									   ALLOCSET_DEFAULT_SIZES);
	fc_oldcontext = MemoryContextSwitchTo(fc_plan->mcxt);

	fc_nargs = fc_list ? PySequence_Length(fc_list) : 0;

	fc_plan->nargs = fc_nargs;
	fc_plan->types = fc_nargs ? palloc0(sizeof(Oid) * fc_nargs) : NULL;
	fc_plan->values = fc_nargs ? palloc0(sizeof(Datum) * fc_nargs) : NULL;
	fc_plan->args = fc_nargs ? palloc0(sizeof(PLyObToDatum) * fc_nargs) : NULL;

	MemoryContextSwitchTo(fc_oldcontext);

	fc_oldcontext = CurrentMemoryContext;
	fc_oldowner = CurrentResourceOwner;

	PLy_spi_subtransaction_begin(fc_oldcontext, fc_oldowner);

	PG_TRY();
	{
		int			fc_i;

		for (fc_i = 0; fc_i < fc_nargs; fc_i++)
		{
			char	   *fc_sptr;
			Oid			fc_typeId;
			int32		fc_typmod;

			fc_optr = PySequence_GetItem(fc_list, fc_i);
			if (PyUnicode_Check(fc_optr))
				fc_sptr = PLyUnicode_AsString(fc_optr);
			else
			{
				ereport(ERROR,
						(errmsg("plpy.prepare: type name at ordinal position %d is not a string", fc_i)));
				fc_sptr = NULL;	/* 保持编译器安静 */
			}

			/********************************************************
			 * 解析参数类型名称，然后通过
			 * oid在系统缓存中查找，并记住
			 * 输入转换所需的信息。
			 ********************************************************/

			parseTypeString(fc_sptr, &fc_typeId, &fc_typmod, false);

			Py_DECREF(fc_optr);

			/*
			 * 将optr设置为NULL，以便在出现
			 * 错误时不会再次尝试取消引用
			 */
			fc_optr = NULL;

			fc_plan->types[fc_i] = fc_typeId;
			PLy_output_setup_func(&fc_plan->args[fc_i], fc_plan->mcxt,
								  fc_typeId, fc_typmod,
								  fc_exec_ctx->curr_proc);
		}

		pg_verifymbstr(fc_query, strlen(fc_query), false);
		fc_plan->plan = SPI_prepare(fc_query, fc_plan->nargs, fc_plan->types);
		if (fc_plan->plan == NULL)
			elog(ERROR, "SPI_prepare failed: %s",
				 SPI_result_code_string(SPI_result));

		/* 将计划从procCxt转移到topCxt */
		if (SPI_keepplan(fc_plan->plan))
			elog(ERROR, "SPI_keepplan failed");

		PLy_spi_subtransaction_commit(fc_oldcontext, fc_oldowner);
	}
	PG_CATCH();
	{
		Py_DECREF(fc_plan);
		Py_XDECREF(fc_optr);

		PLy_spi_subtransaction_abort(fc_oldcontext, fc_oldowner);
		return NULL;
	}
	PG_END_TRY();

	Assert(fc_plan->plan != NULL);
	return (PyObject *) fc_plan;
}

/* execute(query="select * from foo", limit=5)
 * execute(plan=plan, values=(foo, bar), limit=5)
 */
PyObject * PLy_spi_execute(PyObject *fc_self, PyObject *fc_args)
{
	char	   *fc_query;
	PyObject   *fc_plan;
	PyObject   *fc_list = NULL;
	long		fc_limit = 0;

	if (PyArg_ParseTuple(fc_args, "s|l", &fc_query, &fc_limit))
		return fc_PLy_spi_execute_query(fc_query, fc_limit);

	PyErr_Clear();

	if (PyArg_ParseTuple(fc_args, "O|Ol", &fc_plan, &fc_list, &fc_limit) &&
		is_PLyPlanObject(fc_plan))
		return PLy_spi_execute_plan(fc_plan, fc_list, fc_limit);

	PLy_exception_set(PLy_exc_error, "plpy.execute expected a query or a plan");
	return NULL;
}

PyObject * PLy_spi_execute_plan(PyObject *fc_ob, PyObject *fc_list, long fc_limit)
{
	volatile int fc_nargs;
	int			fc_rv;
	PLyPlanObject *fc_plan;
	volatile MemoryContext fc_oldcontext;
	volatile ResourceOwner fc_oldowner;
	PyObject   *fc_ret;

	if (fc_list != NULL)
	{
		if (!PySequence_Check(fc_list) || PyUnicode_Check(fc_list))
		{
			PLy_exception_set(PyExc_TypeError, "plpy.execute takes a sequence as its second argument");
			return NULL;
		}
		fc_nargs = PySequence_Length(fc_list);
	}
	else
		fc_nargs = 0;

	fc_plan = (PLyPlanObject *) fc_ob;

	if (fc_nargs != fc_plan->nargs)
	{
		char	   *fc_sv;
		PyObject   *fc_so = PyObject_Str(fc_list);

		if (!fc_so)
			PLy_elog(ERROR, "could not execute plan");
		fc_sv = PLyUnicode_AsString(fc_so);
		PLy_exception_set_plural(PyExc_TypeError,
								 "Expected sequence of %d argument, got %d: %s",
								 "Expected sequence of %d arguments, got %d: %s",
								 fc_plan->nargs,
								 fc_plan->nargs, fc_nargs, fc_sv);
		Py_DECREF(fc_so);

		return NULL;
	}

	fc_oldcontext = CurrentMemoryContext;
	fc_oldowner = CurrentResourceOwner;

	PLy_spi_subtransaction_begin(fc_oldcontext, fc_oldowner);

	PG_TRY();
	{
		PLyExecutionContext *fc_exec_ctx = PLy_current_execution_context();
		MemoryContext fc_tmpcontext;
		Datum	   *volatile fc_values;
		char	   *volatile fc_nulls;
		volatile int fc_j;

		/*
		 * 转换后的参数和相关冗余信息将存在于这个上下文中，
		 * 该上下文是我们子事务的本地上下文。
		 */
		fc_tmpcontext = AllocSetContextCreate(CurTransactionContext,
										   "PL/Python temporary context",
										   ALLOCSET_SMALL_SIZES);
		MemoryContextSwitchTo(fc_tmpcontext);

		if (fc_nargs > 0)
		{
			fc_values = (Datum *) palloc(fc_nargs * sizeof(Datum));
			fc_nulls = (char *) palloc(fc_nargs * sizeof(char));
		}
		else
		{
			fc_values = NULL;
			fc_nulls = NULL;
		}

		for (fc_j = 0; fc_j < fc_nargs; fc_j++)
		{
			PLyObToDatum *fc_arg = &fc_plan->args[fc_j];
			PyObject   *fc_elem;

			fc_elem = PySequence_GetItem(fc_list, fc_j);
			PG_TRY();
			{
				bool		fc_isnull;

				fc_values[fc_j] = PLy_output_convert(fc_arg, fc_elem, &fc_isnull);
				fc_nulls[fc_j] = fc_isnull ? 'n' : ' ';
			}
			PG_FINALLY();
			{
				Py_DECREF(fc_elem);
			}
			PG_END_TRY();
		}

		MemoryContextSwitchTo(fc_oldcontext);

		fc_rv = SPI_execute_plan(fc_plan->plan, fc_values, fc_nulls,
							  fc_exec_ctx->curr_proc->fn_readonly, fc_limit);
		fc_ret = fc_PLy_spi_execute_fetch_result(SPI_tuptable, SPI_processed, fc_rv);

		MemoryContextDelete(fc_tmpcontext);
		PLy_spi_subtransaction_commit(fc_oldcontext, fc_oldowner);
	}
	PG_CATCH();
	{
		/* 子事务回滚将移除tmpcontext */
		PLy_spi_subtransaction_abort(fc_oldcontext, fc_oldowner);
		return NULL;
	}
	PG_END_TRY();

	if (fc_rv < 0)
	{
		PLy_exception_set(PLy_exc_spi_error,
						  "SPI_execute_plan failed: %s",
						  SPI_result_code_string(fc_rv));
		return NULL;
	}

	return fc_ret;
}

static PyObject * fc_PLy_spi_execute_query(char *fc_query, long fc_limit)
{
	int			fc_rv;
	volatile MemoryContext fc_oldcontext;
	volatile ResourceOwner fc_oldowner;
	PyObject   *fc_ret = NULL;

	fc_oldcontext = CurrentMemoryContext;
	fc_oldowner = CurrentResourceOwner;

	PLy_spi_subtransaction_begin(fc_oldcontext, fc_oldowner);

	PG_TRY();
	{
		PLyExecutionContext *fc_exec_ctx = PLy_current_execution_context();

		pg_verifymbstr(fc_query, strlen(fc_query), false);
		fc_rv = SPI_execute(fc_query, fc_exec_ctx->curr_proc->fn_readonly, fc_limit);
		fc_ret = fc_PLy_spi_execute_fetch_result(SPI_tuptable, SPI_processed, fc_rv);

		PLy_spi_subtransaction_commit(fc_oldcontext, fc_oldowner);
	}
	PG_CATCH();
	{
		PLy_spi_subtransaction_abort(fc_oldcontext, fc_oldowner);
		return NULL;
	}
	PG_END_TRY();

	if (fc_rv < 0)
	{
		Py_XDECREF(fc_ret);
		PLy_exception_set(PLy_exc_spi_error,
						  "SPI_execute failed: %s",
						  SPI_result_code_string(fc_rv));
		return NULL;
	}

	return fc_ret;
}

static PyObject * fc_PLy_spi_execute_fetch_result(SPITupleTable *fc_tuptable, uint64 fc_rows, int fc_status)
{
	PLyResultObject *fc_result;
	PLyExecutionContext *fc_exec_ctx = PLy_current_execution_context();
	volatile MemoryContext fc_oldcontext;

	fc_result = (PLyResultObject *) PLy_result_new();
	if (!fc_result)
	{
		SPI_freetuptable(fc_tuptable);
		return NULL;
	}
	Py_DECREF(fc_result->status);
	fc_result->status = PyLong_FromLong(fc_status);

	if (fc_status > 0 && fc_tuptable == NULL)
	{
		Py_DECREF(fc_result->nrows);
		fc_result->nrows = PyLong_FromUnsignedLongLong(fc_rows);
	}
	else if (fc_status > 0 && fc_tuptable != NULL)
	{
		PLyDatumToOb fc_ininfo;
		MemoryContext fc_cxt;

		Py_DECREF(fc_result->nrows);
		fc_result->nrows = PyLong_FromUnsignedLongLong(fc_rows);

		fc_cxt = AllocSetContextCreate(CurrentMemoryContext,
									"PL/Python temp context",
									ALLOCSET_DEFAULT_SIZES);

		/* 初始化以转换结果元组为Python */
		PLy_input_setup_func(&fc_ininfo, fc_cxt, RECORDOID, -1,
							 fc_exec_ctx->curr_proc);

		fc_oldcontext = CurrentMemoryContext;
		PG_TRY();
		{
			MemoryContext fc_oldcontext2;

			if (fc_rows)
			{
				uint64		fc_i;

				/*
				 * PyList_New()和PyList_SetItem()使用Py_ssize_t作为列表
				 * 大小和列表索引；因此我们无法支持大于
				 * PY_SSIZE_T_MAX的结果。
				 */
				if (fc_rows > (uint64) PY_SSIZE_T_MAX)
					ereport(ERROR,
							(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
							 errmsg("query result has too many rows to fit in a Python list")));

				Py_DECREF(fc_result->rows);
				fc_result->rows = PyList_New(fc_rows);
				if (fc_result->rows)
				{
					PLy_input_setup_tuple(&fc_ininfo, fc_tuptable->tupdesc,
										  fc_exec_ctx->curr_proc);

					for (fc_i = 0; fc_i < fc_rows; fc_i++)
					{
						PyObject   *fc_row = PLy_input_from_tuple(&fc_ininfo,
															   fc_tuptable->vals[fc_i],
															   fc_tuptable->tupdesc,
															   true);

						PyList_SetItem(fc_result->rows, fc_i, fc_row);
					}
				}
			}

			/*
			 * 保存元组描述符以供结果集元数据
			 * 函数稍后使用。 将其保存在TopMemoryContext中，以便它可以在
			 * SPI上下文之外存活。 我们相信PLy_result_dealloc()
			 * 会在适当的时候清理它。 （尽量推迟此操作，
			 * 以最小化由于错误导致tupdesc泄漏的方式。）
			 */
			fc_oldcontext2 = MemoryContextSwitchTo(TopMemoryContext);
			fc_result->tupdesc = CreateTupleDescCopy(fc_tuptable->tupdesc);
			MemoryContextSwitchTo(fc_oldcontext2);
		}
		PG_CATCH();
		{
			MemoryContextSwitchTo(fc_oldcontext);
			MemoryContextDelete(fc_cxt);
			Py_DECREF(fc_result);
			PG_RE_THROW();
		}
		PG_END_TRY();

		MemoryContextDelete(fc_cxt);
		SPI_freetuptable(fc_tuptable);

		/* 如果上面的PyList_New()失败 */
		if (!fc_result->rows)
		{
			Py_DECREF(fc_result);
			fc_result = NULL;
		}
	}

	return (PyObject *) fc_result;
}

PyObject * PLy_commit(PyObject *fc_self, PyObject *fc_args)
{
	MemoryContext fc_oldcontext = CurrentMemoryContext;
	PLyExecutionContext *fc_exec_ctx = PLy_current_execution_context();

	PG_TRY();
	{
		SPI_commit();

		/* 在事务结束时被清除，重置指针 */
		fc_exec_ctx->scratch_ctx = NULL;
	}
	PG_CATCH();
	{
		ErrorData  *fc_edata;
		PLyExceptionEntry *fc_entry;
		PyObject   *fc_exc;

		/* 保存错误信息 */
		MemoryContextSwitchTo(fc_oldcontext);
		fc_edata = CopyErrorData();
		FlushErrorState();

		/* 在事务结束时被清除，重置指针 */
		fc_exec_ctx->scratch_ctx = NULL;

		/* 查找正确的异常 */
		fc_entry = hash_search(PLy_spi_exceptions, &(fc_edata->sqlerrcode),
							HASH_FIND, NULL);

		/*
		 * 这可能是一个自定义错误代码，如果是这样，回退到
		 * SPIError
		 */
		fc_exc = fc_entry ? fc_entry->exc : PLy_exc_spi_error;
		/* 让Python引发异常 */
		fc_PLy_spi_exception_set(fc_exc, fc_edata);
		FreeErrorData(fc_edata);

		return NULL;
	}
	PG_END_TRY();

	Py_RETURN_NONE;
}

PyObject * PLy_rollback(PyObject *fc_self, PyObject *fc_args)
{
	MemoryContext fc_oldcontext = CurrentMemoryContext;
	PLyExecutionContext *fc_exec_ctx = PLy_current_execution_context();

	PG_TRY();
	{
		SPI_rollback();

		/* 在事务结束时被清除，重置指针 */
		fc_exec_ctx->scratch_ctx = NULL;
	}
	PG_CATCH();
	{
		ErrorData  *fc_edata;
		PLyExceptionEntry *fc_entry;
		PyObject   *fc_exc;

		/* 保存错误信息 */
		MemoryContextSwitchTo(fc_oldcontext);
		fc_edata = CopyErrorData();
		FlushErrorState();

		/* 在事务结束时被清除，重置指针 */
		fc_exec_ctx->scratch_ctx = NULL;

		/* 查找正确的异常 */
		fc_entry = hash_search(PLy_spi_exceptions, &(fc_edata->sqlerrcode),
							HASH_FIND, NULL);

		/*
		 * 这可能是一个自定义错误代码，如果是这样，回退到
		 * SPIError
		 */
		fc_exc = fc_entry ? fc_entry->exc : PLy_exc_spi_error;
		/* 让Python引发异常 */
		fc_PLy_spi_exception_set(fc_exc, fc_edata);
		FreeErrorData(fc_edata);

		return NULL;
	}
	PG_END_TRY();

	Py_RETURN_NONE;
}

/*
 * 在子事务中运行SPI函数的实用工具。
 *
 * 用法：
 *
 *	MemoryContext oldcontext = CurrentMemoryContext;
 *	ResourceOwner oldowner = CurrentResourceOwner;
 *
 *	PLy_spi_subtransaction_begin(oldcontext, oldowner);
 *	PG_TRY();
 *	{
 *		<调用SPI函数>
 *		PLy_spi_subtransaction_commit(oldcontext, oldowner);
 *	}
 *	PG_CATCH();
 *	{
 *		<进行清理>
 *		PLy_spi_subtransaction_abort(oldcontext, oldowner);
 *		return NULL;
 *	}
 *	PG_END_TRY();
 *
 * 这些实用工具处理恢复与SPI管理器的连接，并在中止时设置Python异常。
 */
void PLy_spi_subtransaction_begin(MemoryContext fc_oldcontext, ResourceOwner fc_oldowner)
{
	BeginInternalSubTransaction(NULL);
	/* 希望在函数的内存上下文中运行 */
	MemoryContextSwitchTo(fc_oldcontext);
}

void PLy_spi_subtransaction_commit(MemoryContext fc_oldcontext, ResourceOwner fc_oldowner)
{
	/* 提交内部事务，返回外部事务上下文 */
	ReleaseCurrentSubTransaction();
	MemoryContextSwitchTo(fc_oldcontext);
	CurrentResourceOwner = fc_oldowner;
}

void PLy_spi_subtransaction_abort(MemoryContext fc_oldcontext, ResourceOwner fc_oldowner)
{
	ErrorData  *fc_edata;
	PLyExceptionEntry *fc_entry;
	PyObject   *fc_exc;

	/* 保存错误信息 */
	MemoryContextSwitchTo(fc_oldcontext);
	fc_edata = CopyErrorData();
	FlushErrorState();

	/* 中止内部事务 */
	RollbackAndReleaseCurrentSubTransaction();
	MemoryContextSwitchTo(fc_oldcontext);
	CurrentResourceOwner = fc_oldowner;

	/* 查找正确的异常 */
	fc_entry = hash_search(PLy_spi_exceptions, &(fc_edata->sqlerrcode),
						HASH_FIND, NULL);

	/*
	 * 这可能是一个自定义错误代码，如果是这种情况则回退到
	 * SPIError
	 */
	fc_exc = fc_entry ? fc_entry->exc : PLy_exc_spi_error;
	/* 让Python引发异常 */
	fc_PLy_spi_exception_set(fc_exc, fc_edata);
	FreeErrorData(fc_edata);
}

/*
 * 引发一个 SPIError，传入更多的错误细节，如
 * 内部查询和错误位置。
 */
static void fc_PLy_spi_exception_set(PyObject *fc_excclass, ErrorData *fc_edata)
{
	PyObject   *fc_args = NULL;
	PyObject   *fc_spierror = NULL;
	PyObject   *fc_spidata = NULL;

	fc_args = Py_BuildValue("(s)", fc_edata->message);
	if (!fc_args)
		goto failure;

	/* 创建一个新的 SPI 异常，将错误消息作为参数 */
	fc_spierror = PyObject_CallObject(fc_excclass, fc_args);
	if (!fc_spierror)
		goto failure;

	fc_spidata = Py_BuildValue("(izzzizzzzz)", fc_edata->sqlerrcode, fc_edata->detail, fc_edata->hint,
							fc_edata->internalquery, fc_edata->internalpos,
							fc_edata->schema_name, fc_edata->table_name, fc_edata->column_name,
							fc_edata->datatype_name, fc_edata->constraint_name);
	if (!fc_spidata)
		goto failure;

	if (PyObject_SetAttrString(fc_spierror, "spidata", fc_spidata) == -1)
		goto failure;

	PyErr_SetObject(fc_excclass, fc_spierror);

	Py_DECREF(fc_args);
	Py_DECREF(fc_spierror);
	Py_DECREF(fc_spidata);
	return;

failure:
	Py_XDECREF(fc_args);
	Py_XDECREF(fc_spierror);
	Py_XDECREF(fc_spidata);
	elog(ERROR, "could not convert SPI error to Python exception");
}
